import numpy as np
import matplotlib.pyplot as plt

# k近邻算法类
class KJinLin:
    def __init__(self, k=3):
        self.k = k  # 选几个邻居
        self.shu_ju = None  # 存训练数据
        self.biao_qian = None  # 存训练标签

    # 训练函数（其实就是存数据）
    def xun_lian(self, xun_lian_shu_ju, xun_lian_biao_qian):
        self.shu_ju = xun_lian_shu_ju
        self.biao_qian = xun_lian_biao_qian

    # 预测函数
    def yu_ce(self, ce_shi_shu_ju):
        yu_ce_jie_guo = []
        # 逐个预测测试数据
        for ce_shi_dian in ce_shi_shu_ju:
            # 算这个测试点到所有训练点的距离
            ju_li_list = []
            for xun_lian_dian in self.shu_ju:
                # 欧氏距离，一步一步算
                ju_li = 0
                for i in range(len(ce_shi_dian)):
                    ju_li = ju_li + (ce_shi_dian[i] - xun_lian_dian[i]) **2
                ju_li = np.sqrt(ju_li)
                ju_li_list.append(ju_li)
            
            # 找最近的k个邻居的索引
            zui_jin_de = np.argsort(ju_li_list)[:self.k]
            
            # 统计这k个邻居的标签
            biao_qian_tong_ji = {}
            for idx in zui_jin_de:
                bq = self.biao_qian[idx]
                if bq in biao_qian_tong_ji:
                    biao_qian_tong_ji[bq] = biao_qian_tong_ji[bq] + 1
                else:
                    biao_qian_tong_ji[bq] = 1
            
            # 找出现次数最多的标签
            zui_duo = -1
            zui_hou_bq = 0
            for key, value in biao_qian_tong_ji.items():
                if value > zui_duo:
                    zui_duo = value
                    zui_hou_bq = key
            yu_ce_jie_guo.append(zui_hou_bq)
        
        return np.array(yu_ce_jie_guo)


# 测试一下
if __name__ == "__main__":
    # 自己造点二维数据，分两类
    def zao_shu_ju():
        # 第一类数据，标签0
        lei0 = np.random.normal(2, 1, (50, 2))
        bq0 = np.zeros(50)
        # 第二类数据，标签1
        lei1 = np.random.normal(7, 1, (50, 2))
        bq1 = np.ones(50)
        # 合并
        all_shu_ju = np.vstack((lei0, lei1))
        all_biao_qian = np.hstack((bq0, bq1))
        return all_shu_ju, all_biao_qian

    # 生成训练数据
    xun_lian_shu_ju, xun_lian_biao_qian = zao_shu_ju()

    # 建个模型，k=5
    model = KJinLin(k=5)
    model.xun_lian(xun_lian_shu_ju, xun_lian_biao_qian)

    # 造几个测试点
    ce_shi_shu_ju = np.array([[2, 2], [3, 3], [6, 7], [8, 8], [5, 5]])
    # 预测
    ce_shi_jie_guo = model.yu_ce(ce_shi_shu_ju)
    print("测试点的预测结果：", ce_shi_jie_guo)

    # 画图
    plt.scatter(xun_lian_shu_ju[:, 0], xun_lian_shu_ju[:, 1], c=xun_lian_biao_qian, cmap='coolwarm', s=30, label='训练数据')
    plt.scatter(ce_shi_shu_ju[:, 0], ce_shi_shu_ju[:, 1], c=ce_shi_jie_guo, cmap='coolwarm', s=100, marker='*', label='测试数据')
    plt.title('k近邻分类结果')
    plt.legend()
    plt.show()